/*

 * Copyright (c) 2017-2018, FastDev 刘强 (fastdev@163.com).

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package of.hotswapinject.injection;

import of.hotswapinject.core.common.Const;
import of.hotswapinject.core.common.Util;
import org.objectweb.asm.*;
import of.hotswapinject.core.plugin.IMethodFilter;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.List;

public class MethodHookTransformer implements ClassFileTransformer {
    private String savePath;
    private HashSet<String> matchedClasses = null;
    private IMethodFilter methodFilter = null;

    public MethodHookTransformer(IMethodFilter methodFilter, List<String> matchedClasses,
                             String savePath) {
        this.methodFilter = methodFilter;
        this.matchedClasses = new HashSet<String>();
        for(String str : matchedClasses) {
            this.matchedClasses.add(str.replace('.', '/'));
        }
        this.savePath = savePath;
    }

    public boolean isMatch(String name) {
        return this.matchedClasses.contains(name);
    }

    private ClassWriter getClassWriter(String className, byte[] classfileBuffer, IMethodHook hookMethod) {
        ClassReader reader = new ClassReader(classfileBuffer);
        ClassWriter writer = new ClassWriter(reader, hookMethod.computeFrames() ? ClassWriter.COMPUTE_FRAMES : ClassWriter.COMPUTE_MAXS);
        ClassVisitor visitor = new MethodHookClassVisitor(hookMethod, writer);
        reader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return writer;
    }

    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(!isMatch(className)) {
            return null;
        }

        try {
            MethodHookedChecker hookedChecker = new MethodHookedChecker();
            getClassWriter(className, classfileBuffer, hookedChecker);
            if(hookedChecker.isHooked()) {
                Util.trace("Class:" + className + " has already been hooked!");
            }
            MethodHook hookMethod = new MethodHook(methodFilter, className, hookedChecker.isHooked());
            ClassWriter writer = getClassWriter(className, classfileBuffer, hookMethod);
            byte[] result = writer.toByteArray();
            if (savePath != null && !("".equals(savePath))) {
                if(!savePath.endsWith(Const.FILESEPARATOR)) {
                    savePath = savePath + Const.FILESEPARATOR;
                }
                Util.saveFile(savePath + className + ".class", result);
            }
            return result;
        }
        catch(Throwable e) {
            e.printStackTrace();
            return null;
        }
    }
}
